import numpy as np
import torch
import torch.nn as nn
from PIL import Image
import matplotlib.pyplot as plt
from torchvision import transforms

"""
图像预处理：读取图片转化为灰度图，并转化为Numpy数组,也可以直接转为tensor
"""
img = Image.open('../images/dog.png')
# 直接使用 torchvision.transforms.ToTensor() 将 PIL 图像转换为 PyTorch 张量是更简洁且推荐的做法。
# img_gray = transforms.ToTensor()(img.convert('L'))
# 将图片转化为灰度图，并转化为Numpy数组，方便后续进行操作
img_gray = np.array(img.convert('L'), dtype=np.float32)
plt.figure(figsize=(6, 6))
# cmap='gray'：指定颜色映射为灰度（'gray' 是 plt.cm.gray 的简写）。
plt.imshow(img_gray, cmap='gray')
# plt.axis('off')：关闭坐标轴显示（让图像更干净）。
plt.axis('off')
plt.show()

"""
图像预处理：将图片转化为张量
"""
# print(img_gray.shape)
imh, imw = img_gray.shape
"""
PyTorch 的卷积操作（如 nn.Conv2d）要求输入张量的维度为：(batch_size, channels, height, width)
作用：将 (imh, imw) 的二维数组转换为 四维张量，形状变为 (1, 1, imh, imw)。
维度含义：
    第1维 1：批处理维度（batch size），表示输入1张图片。
    第2维 1：通道数（channel），灰度图是单通道。
    第3~4维 imh, imw：图像的高度和宽度。
等价于：
    from torchvision import transforms

    img_tensor = transforms.ToTensor()(img_gray)  # 自动转为 [C, H, W]（这里是 [1, 256, 256]）
    img_tensor = img_tensor.unsqueeze(0)          # 添加batch维度 → [1, 1, 256, 256]
"""
img_tensor = torch.from_numpy(img_gray.reshape(1, 1, imh, imw))

"""
使用5*5的随机数构成的卷积核进行卷积操作
"""
conv_kernel = torch.randn(1, 1, 5, 5)
"""
nn.functional.conv2d 的参数灵活控制了卷积的计算方式和输出尺寸
实际应用中，通常直接使用 nn.Conv2d 模块（封装了 conv2d），但底层原理完全一致。
"""
# conv_out = nn.functional.conv2d(img_tensor, conv_kernel)
conv_out = nn.Conv2d(1, 1, 5, 5)(img_tensor)
print(conv_out.shape)
"""
可视化卷积结果：
    squeeze()：去掉张量中维度为1的轴（从 [1, 1, 252, 252] 变为 [252, 252]）。
    .numpy()：将 PyTorch 张量转为 NumPy 数组供 Matplotlib 显示。
    cmap='gray'：以灰度图显示。
    效果：显示卷积后的特征图，可能呈现边缘、纹理等局部特征。
"""
plt.figure(figsize=(6, 6))
"""
conv_out.squeeze()       # 移除维度为1的轴
plt.imshow() 只接受2D（灰度）或3D（RGB）数组，需去掉多余的batch和channel维度。
"""
plt.imshow(conv_out.squeeze().detach().numpy(), cmap='gray') # 采用nn.Conv2d的方式获取结果
# plt.imshow(conv_out.squeeze().numpy(), cmap='gray')
plt.axis('off')
plt.show()
print(torch.__version__)